/**
 * 
 */
package edu.cmu.km.kernel.tree;

import java.util.List;

import edu.cmu.km.core.*;

/**
 * @author Lucas Xu
 * @param <T>
 * 
 */
public class TreeKernel<T> {

	/**
	 * Co-rooted subtree kernel Reference: Kernel Methods for Pattern Analysis,
	 * Code Fragment 11.7.Pseudocode for the co-rooted subtree kernel.
	 * 
	 * @param tree1
	 * @param tree2
	 * @return
	 */
	public int corootedSubtreeKernel(Tree<T> tree1, Tree<T> tree2) {
		return corootedSubtreeProcessnode(tree1.getRootElement(), tree2
				.getRootElement());
	}

	private int corootedSubtreeProcessnode(Node<T> rootTree1, Node<T> rootTree2) {
		int num_children = rootTree1.getNumberOfChildren();

		if (num_children != rootTree2.getNumberOfChildren()
				|| num_children == 0)
			return 0;
		else {
			int kern = 1;
			for (int i = 0; i < num_children; i++) {
				kern = kern
						* (corootedSubtreeProcessnode(rootTree1.getChildren()
								.get(i), rootTree2.getChildren().get(i)) + 1);

			}
			return kern;
		}
	}

	/**
	 * All subtree kernel, Reference: Kernel Methods for Pattern Analysis, Code
	 * Fragment 11.8.Code Fragment 11.8. Pseudocode for the all-subtree kernel.
	 * Implemented in dynamic programming
	 * 
	 * @param tree1
	 * @param tree2
	 * @return
	 */
	public int allSubtreeKernel(Tree<T> tree1, Tree<T> tree2) {
		// put all nodes in the tree in a post-order list so that parents are
		// always after the children.
		List<Node<T>> nodes1 = tree1.postTraverse();
		List<Node<T>> nodes2 = tree2.postTraverse();
		int n1 = nodes1.size();
		int n2 = nodes2.size();

		if (n1 <= 1 || n2 <= 1)
			return 0;

		// calculate the co-rooted subtree kernel dynamic programming matrix
		int[][] DPr = new int[n1][n2];
		System.out.println();
		for (int i = 0; i < n1; i++) {
			for (int j = 0; j < n2; j++) {
				List<Node<T>> childrenOfNodei = nodes1.get(i).getChildren();
				List<Node<T>> childrenOfNodej = nodes2.get(j).getChildren();
				int[] childrenIndexOfNodei = new int[childrenOfNodei.size()];
				int[] childrenIndexOfNodej = new int[childrenOfNodej.size()];
				for (int m = 0; m < childrenIndexOfNodei.length; m++) {
					childrenIndexOfNodei[m] = nodes1.indexOf(childrenOfNodei
							.get(m));
				}
				for (int n = 0; n < childrenIndexOfNodej.length; n++) {
					childrenIndexOfNodej[n] = nodes2.indexOf(childrenOfNodej
							.get(n));
				}

				int num_children = nodes1.get(i).getNumberOfChildren();
				if (num_children != nodes2.get(j).getNumberOfChildren()
						|| num_children == 0)
					DPr[i][j] = 0;
				else {
					DPr[i][j] = 1;
					for (int k = 0; k < num_children; k++) {
						DPr[i][j] = DPr[i][j]
								* (DPr[childrenIndexOfNodei[k]][childrenIndexOfNodej[k]] + 1);
					}
				}
				System.out.print(" " + DPr[i][j]);
			}
			System.out.println();
		}

		// calculate the all subtree kernels recursively
		int[][] DP = new int[n1][n2];
		for (int i = 0; i < n1; i++)
			for (int j = 0; j < n2; j++) {
				List<Node<T>> childrenOfNodei = nodes1.get(i).getChildren();
				List<Node<T>> childrenOfNodej = nodes2.get(j).getChildren();
				int[] childrenIndexOfNodei = new int[childrenOfNodei.size()];
				int[] childrenIndexOfNodej = new int[childrenOfNodej.size()];
				for (int m = 0; m < childrenIndexOfNodei.length; m++) {
					childrenIndexOfNodei[m] = nodes1.indexOf(childrenOfNodei
							.get(m));
				}
				for (int n = 0; n < childrenIndexOfNodej.length; n++) {
					childrenIndexOfNodej[n] = nodes1.indexOf(childrenOfNodej
							.get(n));
				}

				if (nodes1.get(i).getNumberOfChildren() == 0
						|| nodes2.get(j).getNumberOfChildren() == 0)
					DP[i][j] = 0;
				else {
					DP[i][j] = DPr[i][j];
					for (int m = 0; m < childrenIndexOfNodei.length; m++) {
						DPr[i][j] += DP[childrenIndexOfNodei[m]][j];
					}
					for (int n = 0; n < childrenIndexOfNodej.length; n++) {
						DPr[i][j] += DPr[i][childrenIndexOfNodej[n]];
					}
					for (int m = 0; m < childrenIndexOfNodei.length; m++)
						for (int n = 0; n < childrenIndexOfNodej.length; n++) {
							DPr[i][j] -= DPr[childrenIndexOfNodei[m]][childrenIndexOfNodej[n]];
						}
				}
			}

		return DP[n1 - 1][n2 - 1];
	}

	/**
	 * size-of-agtreement subtree distribution kernel, for ordered and labeled
	 * tree Reference: Kernel Based on Distributions of Agreement Subtrees by
	 * Kilho Shin, and Tetsuji Kuboyama k(x,y)=sum(f(sizeof(t))), where t is the
	 * agreement subtree of (x,y) f=alpha^sizeof(t)
	 * 
	 * @param tree1
	 *            ordered and labeled tree
	 * @param tree2
	 *            ordered and labeled tree
	 * @return
	 */
	public int sizeOfASTKernel_alpha(Tree<T> tree1, Tree<T> tree2) {

		return 0;

	}

	public int sparseSubForestKernel(Tree<T> tree1, Tree<T> tree2) {
		return this.sparseSubForestKernel(tree1.toForest(), tree2.toForest());
	}

	/**
	 * Sparse SubForest kernel utility
	 * 
	 * @param forest_x
	 * @param forest_y
	 * @return
	 */
	public int sparseSubForestKernel(Forest<T> forest_x, Forest<T> forest_y) {
		int kernel = 0;

		// if forest_x is empty or forest_y is empty, return 0;
		if (forest_x.getFirstRootElement() == null
				|| forest_y.getFirstRootElement() == null)
			kernel = 0;
		else {
			// co-rooted sparse subforest kernel *K(X,Y)
			if (forest_x.getFirstRootElement().getData() == forest_y
					.getFirstRootElement().getData()) {
				if (forest_x.getForestWithoutFirstRoot().getFirstRootElement() == null
						&& forest_y.getForestWithoutFirstRoot()
								.getFirstRootElement() == null)
					kernel = 1;
				else {
					int a1 = this.sparseSubForestKernel(forest_x
							.getForestWithoutMaximumCorootedSubtree(), forest_y
							.getForestWithoutMaximumCorootedSubtree());
					int a2 = this.sparseSubForestKernel(forest_x
							.getMaximumCorootedSubtree().toForest()
							.getForestWithoutFirstRoot(), forest_y
							.getMaximumCorootedSubtree().toForest()
							.getForestWithoutFirstRoot());
					kernel += (1 + a1) * (1 + a2);
				}
			}

			int b = this.sparseSubForestKernel(forest_x, forest_y
					.getForestWithoutFirstRoot());
			int c = this.sparseSubForestKernel(forest_x
					.getForestWithoutFirstRoot(), forest_y);
			int d = this.sparseSubForestKernel(forest_x
					.getForestWithoutFirstRoot(), forest_y
					.getForestWithoutFirstRoot());

			kernel += b + c - d;
		}
		System.out.print("forest_x: " + forest_x.toString());
		System.out.print("forest_y: " + forest_y.toString());

		System.out.println("Kernel: " + kernel);
		return kernel;

	}

	private int phi(Node<T> x, Node<T> y) {
		if (x.data == y.data)
			return 1;
		else
			return 0;
	}

	private int pi(Node<T> x, Node<T> y) {
		if (x.data == y.data)
			return 1;
		else
			return 0;
	}

	/**
	 * Sparse subtree, agreement subtree, Contiguous, and continuous subtree
	 * kernel.
	 * 
	 * @param forest_x
	 * @param forest_y
	 * @param n
	 * @return
	 */
	public int SACCoTreeKernel(Forest<T> forest_x, Forest<T> forest_y, int n) {
		int kernel = 0;
		// if forest_x is empty or forest_y is empty, return 0;
		if (forest_x.getFirstRootElement() == null
				|| forest_y.getFirstRootElement() == null)
			kernel = 0;
		else {
			int phi = this.phi(forest_x.getFirstRootElement(), forest_y
					.getFirstRootElement());
			int pi = this.pi(forest_x.getFirstRootElement(), forest_y
					.getFirstRootElement());

			int a = (int) (phi * Math.pow(pi, n));
			// initialize C(n,i)
			int[] cni = new int[n + 1];
			cni[0] = 1;
			cni[n] = 1;
			for (int i = 1; i <= n / 2; i++) {
				cni[i] = (n - i + 1) * cni[i - 1] / i;
				cni[n - i] = cni[i];
			}
			for (int i = 0; i <= n; i++) {
				double recur = this.SACCoTreeKernel(forest_x
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), forest_y
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), n - i);
				if (recur == 0)
					recur = 1;
				a += cni[i] * phi * Math.pow(pi, i) * recur;
			}

			int b = this.SACCoTreeKernel(forest_x.getMaximumCorootedSubtree()
					.toForest(), forest_y.getMaximumCorootedSubtree()
					.toForest().getForestWithoutFirstRoot(), n);
			int c = this.SACCoTreeKernel(forest_x.getMaximumCorootedSubtree()
					.toForest().getForestWithoutFirstRoot(), forest_y
					.getMaximumCorootedSubtree().toForest(), n);
			int d = this.SACCoTreeKernel(forest_x.getMaximumCorootedSubtree()
					.toForest().getForestWithoutFirstRoot(), forest_y
					.getMaximumCorootedSubtree().toForest()
					.getForestWithoutFirstRoot(), n);

			int e = this.SACCoTreeKernel(forest_x, forest_y
					.getForestWithoutMaximumCorootedSubtree(), n);
			int f = this.SACCoTreeKernel(forest_x
					.getForestWithoutMaximumCorootedSubtree(), forest_y, n);
			int g = this.SACCoTreeKernel(forest_x
					.getForestWithoutMaximumCorootedSubtree(), forest_y
					.getForestWithoutMaximumCorootedSubtree(), n);

			kernel = a + b + c - d + e + f - g;
		}

		return kernel;
	}

	public int SACCoTreeKernel(Tree<T> tree1, Tree<T> tree2, int n) {
		return this.SACCoTreeKernel(tree1.toForest(), tree2.toForest(), n);
	}

	/**
	 * agreement sub forest kernel, $=AF
	 * 
	 * @param forest_x
	 * @param forest_y
	 * @param n
	 * @return
	 */
	public int agreementSubForestKernel(Forest<T> forest_x, Forest<T> forest_y,
			int n) {
		int kernel = 0;
		// if forest_x is empty or forest_y is empty, return 0;
		if (forest_x.getFirstRootElement() == null
				|| forest_y.getFirstRootElement() == null)
			kernel = 0;
		else {
			int phi = this.phi(forest_x.getFirstRootElement(), forest_y
					.getFirstRootElement());
			int pi = this.pi(forest_x.getFirstRootElement(), forest_y
					.getFirstRootElement());
			int a = 0;

			// initialize C(n,i)
			int[] cni = new int[n + 1];
			cni[0] = 1;
			cni[n] = 1;
			for (int i = 1; i <= n / 2; i++) {
				cni[i] = (n - i + 1) * cni[i - 1] / i;
				cni[n - i] = cni[i];
			}
			for (int i = 0; i <= n; i++) {
				double recur = this.agreementSubForestKernel(forest_x
						.getForestWithoutMaximumCorootedSubtree(), forest_y
						.getForestWithoutMaximumCorootedSubtree(), n - i);
				if (recur == 0)
					recur = 1;

				a += cni[i] * phi * Math.pow(pi, i) * recur;
			}

			// C(n,i,j)
			int[][] cnij = new int[n + 1][n + 1];
			for (int i = 0; i <= n; i++) {
				cnij[i][0] = cni[i];
				for (int j = 1; j <= n - i; j++) {
					cnij[i][j] = cnij[i][j - 1] * (n - i - j + 1) / j;
				}
			}

			for (int i = 0; i <= n; i++) {
				for (int j = 0; j <= n - i; j++) {
					double recurj = this.agreementSubForestKernel(forest_x
							.getMaximumCorootedSubtree().toForest()
							.getForestWithoutFirstRoot(), forest_y
							.getMaximumCorootedSubtree().toForest()
							.getForestWithoutFirstRoot(), j);
					double recurk = this.agreementSubForestKernel(forest_x
							.getForestWithoutMaximumCorootedSubtree(), forest_y
							.getForestWithoutMaximumCorootedSubtree(), n - i
							- j);
					// incase the K(AF, i) is 0
					if (recurj == 0)
						recurj = 1;
					if (recurk == 0)
						recurk = 1;
					a += cnij[i][j] * phi * Math.pow(pi, i) * recurj * recurk;
				}
			}

			int b = 0;
			for (int i = 0; i <= n; i++) {
				double recuri = this.agreementSubForestKernel(forest_x
						.getMaximumCorootedSubtree().toForest(), forest_y
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), i);
				if (recuri == 0)
					recuri = 1;
				double recurj = this.agreementSubForestKernel(forest_x
						.getForestWithoutMaximumCorootedSubtree(), forest_y
						.getForestWithoutMaximumCorootedSubtree(), n - i);
				if (recurj == 0)
					recurj = 1;
				b += cni[i] * recuri * recurj;
			}
			int c = 0;
			for (int i = 0; i <= n; i++) {
				double recuri = this.agreementSubForestKernel(forest_x
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), forest_y
						.getMaximumCorootedSubtree().toForest(), i);
				if (recuri == 0)
					recuri = 1;
				double recurj = this.agreementSubForestKernel(forest_x
						.getForestWithoutMaximumCorootedSubtree(), forest_y
						.getForestWithoutMaximumCorootedSubtree(), n - i);
				if (recurj == 0)
					recurj = 1;
				c += cni[i] * recuri * recurj;
			}
			int d = 0;
			for (int i = 0; i <= n; i++) {
				double recuri = this.agreementSubForestKernel(forest_x
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), forest_y
						.getMaximumCorootedSubtree().toForest()
						.getForestWithoutFirstRoot(), i);
				if (recuri == 0)
					recuri = 1;
				double recurj = this.agreementSubForestKernel(forest_x
						.getForestWithoutMaximumCorootedSubtree(), forest_y
						.getForestWithoutMaximumCorootedSubtree(), n - i);
				if (recurj == 0)
					recurj = 1;
				d += cni[i] * recuri * recurj;
			}

			int e = this.agreementSubForestKernel(forest_x, forest_y
					.getForestWithoutMaximumCorootedSubtree(), n);
			int f = this.agreementSubForestKernel(forest_x
					.getForestWithoutMaximumCorootedSubtree(), forest_y, n);
			int g = this.agreementSubForestKernel(forest_x
					.getForestWithoutMaximumCorootedSubtree(), forest_y
					.getForestWithoutMaximumCorootedSubtree(), n);

			kernel = a + b + c - d + e + f - g;
		}

		return kernel;

	}
	
	public int agreementSubForestKernel(Tree<T> tree1, Tree<T> tree2, int n) {
		return this.agreementSubForestKernel(tree1.toForest(), tree2.toForest(), n);
	}

}
